import tensorflow as tf

# 指定TensorFlow使用GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 设置TensorFlow增长式内存分配
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 构建一个简单的神经网络模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 加载并预处理数据（这里省略具体实现）
# ...

# 训练模型
model.fit(x_train, y_train, epochs=10)
